package org.laizili.solution.leetcode;

import java.util.Comparator;
import java.util.PriorityQueue;

/**
 * <a href="https://leetcode.cn/problems/kth-smallest-element-in-a-sorted-matrix/">378. 有序矩阵中第 K 小的元素</a>
 * <p>
 *     tags: 二分搜索 + 计数; 多路归并排序（堆维护）;
 */
public class Problem378 {
    private static class Solution {
        public int kthSmallest(int[][] matrix, int k) {
            PriorityQueue<int[]> pq = new PriorityQueue<>(Comparator.comparingInt(l -> l[0]));
            for (int i = 0; i < matrix.length; i++) {
                pq.offer(new int[]{matrix[i][0], i, 0});
            }
            int col = matrix.length;
            for (int i = 0; i < k - 1; i++) {
                int[] head = pq.poll();
                if (head[2] < col - 1) {
                    head[2] += 1;
                    head[0] = matrix[head[1]][head[2]];
                    pq.offer(head);
                }
            }
            return pq.peek()[0];
        }
    }

    public static void main(String[] args) {
        System.out.println(new Solution().kthSmallest(new int[][]{{1, 5, 9}, {10, 11, 13}, {12, 13, 15}
        }, 3));
    }
}
